from mmseg.models.decode_heads.decode_head import BaseDecodeHead
import torch
from torch import nn
from mmseg.models.utils import resize
from mmcv.cnn import ConvModule
from ..backbone.mit import mit_b2
from torch.nn import functional as F
from mmengine import MODELS
from mmengine import init_default_scope
from ..head.segformer_head import SegformerHead

import yaml 

def dropout(xs):
    dropout=nn.Dropout2d(0.5)
    if isinstance(xs, list):
        out = []
        for _x in xs:
            out.append(dropout(_x))
        return out
    else:
        return dropout(xs)
class Segformer(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.backbone = mit_b2()
        self.head = SegformerHead(
            **cfg
        )
        if cfg['pretrained']:
            print("loding weights")
            load_weights_and_check_compatibility(self.backbone, cfg['pretrained'])
            
        with open('model.txt',  'w+') as f:
            print(self, file=f)
    def forward(self, x, use_fp=False):
        h, w = x.shape[-2:]
        x = self.backbone.forward(x)
        if use_fp:
            x = [torch.cat((xi, dropout(xi))) for xi in x]
            out = self.head.forward(x)
            out = F.interpolate(out, size=(h, w), mode="bilinear", align_corners=True)
            return out.chunk(2)
        x = self.head.forward(x)
        out = F.interpolate(x, size=(h, w), mode="bilinear", align_corners=True)
        return out
def load_weights_and_check_compatibility(model, weights_path):
    checkpoint = torch.load(weights_path)
    model_state_dict = model.state_dict()
    
    incompatible_weights = []
    
    for name, param in checkpoint.items():
        if name in model_state_dict:
            if param.shape != model_state_dict[name].shape:
                incompatible_weights.append(name)
        else:
            incompatible_weights.append(name)
    print('incompatible_weights:\n', incompatible_weights)
    return incompatible_weights

def get_segformer(cfg):
    init_default_scope('mmseg')
    model = MODELS.build(cfg)
    print(f'model is {model}')
    return model


def main():
    with open('configs/pascal_segformer.yaml') as f:
        cfg = yaml.load(f)
    model = Segformer(cfg)
    # with open('model')
    # print(model )
    
    x = torch.rand(2, 3, 224, 224)
    out = model(x)
    print(out.shape)
    
if __name__ == '__main__':
    main()
    